import numpy as np
from keras.callbacks import Callback
from keras.utils import to_categorical
from keras.models import load_model
from config.constant import NETWORK

class DLSCACallback(Callback):
    def __init__(self, attacker):
        self.attacker = attacker
        super().__init__()

    def on_epoch_end(self, epoch, logs=None):
        self.attacker.accuracy.append(logs['accuracy'])
        self.attacker.loss.append(logs['loss'])
        if 'val_loss' in logs:
            self.attacker.val_loss.append(logs['val_loss'])
            self.attacker.val_accuracy.append(logs['val_accuracy'])

    def on_batch_end(self, batch, logs=None):
        if self.attacker.stop_train:
            self.model.stop_training = True


def profile(self):
    if len(self.midvalue_profile.shape) == 2:
        self.output_class = self.midvalue_profile.shape[1]
    else:
        max_mv=np.max(self.midvalue_profile)
        if max_mv>100:
            self.output_class=256
        elif max_mv>2:
            self.output_class = 9
        else:
            self.output_class = 2
    if self.continue_train:
        self.model = load_model(self.continue_train_path)
    else:
        model_func = NETWORK[self.param_dict["network"]]
        self.model = model_func(self)
    train_model_callback = DLSCACallback(self)
    if isinstance(self.model.get_layer(index=0).input_shape, list):
        input_layer_shape = self.model.get_layer(index=0).input_shape[0]
    else:
        input_layer_shape = self.model.get_layer(index=0).input_shape
    train_traces_reshape = self.traces_profile
    if len(input_layer_shape) == 3:
        train_traces_reshape = self.traces_profile.reshape((self.traces_profile.shape[0], self.traces_profile.shape[1], 1))
    train_labels = self.midvalue_profile
    if len(self.midvalue_profile.shape) == 1:
        train_labels = to_categorical(self.midvalue_profile, num_classes=self.output_class)
    self.model.fit(x=train_traces_reshape, y=train_labels, batch_size=self.param_dict["batch_size"], verbose=0, epochs=self.param_dict["epochs"], callbacks=[train_model_callback], validation_split=self.param_dict["validation_split"])


class ProbCalculator():
    def __init__(self, model):
        self.model = model
        self.guess_key_likelihood = np.zeros((256))

    def update(self, traces_attack, midvalue):
        predict_matrix = self.model.predict(traces_attack, verbose=0)
        if len(midvalue.shape) == 2:
            likelihood_matrix = np.array([predict_matrix[i, midvalue[i]] for i in range(predict_matrix.shape[0])])
        else:
            likelihood_matrix = np.prod(2 * predict_matrix[:, None, :] * midvalue - predict_matrix[:, None, :] - midvalue + 1, axis=2)
        min_nzero = min(likelihood_matrix[likelihood_matrix != 0])
        likelihood_matrix = np.where(likelihood_matrix == 0, min_nzero, likelihood_matrix)
        guess_key_likelihood = np.sum(np.log(likelihood_matrix), axis=0)
        self.guess_key_likelihood = self.guess_key_likelihood + guess_key_likelihood

    def get_prob(self):
        return self.guess_key_likelihood


def attack(self):
    trs_num = self.traces_attack.shape[0]
    attack_steps = self.param_dict["attack_steps"]
    index = [i for i in range(0, trs_num, attack_steps)]
    index.append(trs_num)
    if "model" not in self.__dict__:
        self.model=load_model(self.param_dict["DLSCA_model_path"])
    if isinstance(self.model.get_layer(index=0).input_shape, list):
        input_layer_shape = self.model.get_layer(index=0).input_shape[0]
    else:
        input_layer_shape = self.model.get_layer(index=0).input_shape
    traces_attack = self.traces_attack
    if len(input_layer_shape) == 3:
        traces_attack = traces_attack.reshape((traces_attack.shape[0], traces_attack.shape[1], 1))
    prob_calculator = ProbCalculator(self.model)
    for i in range(len(index) - 1):
        prob_calculator.update(traces_attack[index[i]:index[i + 1]], self.midvalue_attack[index[i]:index[i + 1]])
        prob = prob_calculator.get_prob()
        yield prob, index[i + 1]


def get_guess_key(self, prob):
    return np.argmax(prob)


def get_rank(self, prob, real_key):
    prob_sort = np.argsort(np.argsort(-prob))
    return prob_sort[real_key]


def get_accuracy(self):
    target_byte = self.param_dict["target_byte"]
    mkey_byte = self.mkey[target_byte]
    real_midvalue = self.midvalue_attack[:, mkey_byte]
    traces_num = self.traces_attack.shape[0]
    if isinstance(self.model.get_layer(index=0).input_shape, list):
        input_layer_shape = self.model.get_layer(index=0).input_shape[0]
    else:
        input_layer_shape = self.model.get_layer(index=0).input_shape
    traces_attack = self.traces_attack
    if len(input_layer_shape) == 3:
        traces_attack = traces_attack.reshape((traces_attack.shape[0], traces_attack.shape[1], 1))
    predict_matrix = self.model.predict(traces_attack, verbose=0)
    if len(real_midvalue) == 1:
        predict_midvlue = np.argmax(predict_matrix, axis=1)
    else:
        vec = 2 ** np.arange(real_midvalue.shape[1])[::-1]
        real_midvalue = real_midvalue@vec
        predict_midvlue = predict_matrix.round()@vec
    accuracy = ((real_midvalue == predict_midvlue).sum()) / traces_num
    return accuracy
